# -*- coding: utf-8 -*-
"""
    Powered by AIgames of michine intellengence lab
    Author: Justin Xu
"""
#Convert the go's numpy to TFRecord

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import sys
import numpy as np

import tensorflow as tf


root_dir = '/home/milab/workspace/generateTest/'
def _int64_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))

def _bytes_feature(value):
     return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

lst_for_go = open("total_data.lst")
filename = os.path.join(root_dir, 'validation' + '.tfrecords')
writer = tf.python_io.TFRecordWriter(filename)
index = 0
while True:
    line = lst_for_go.readline()
    if not line:
        break
    _, data_img_name, label_img_name = line.strip('\n').split("\t")
    images = np.load(root_dir+data_img_name)
    labels = np.load(root_dir+label_img_name)
    labels = labels.reshape((361,))
    # import pdb; pdb.set_trace()
    rows = images.shape[1]
    cols = images.shape[2]
    depth = images.shape[0]
    images_t = np.zeros((19,19,2))
    # import pdb; pdb.set_trace()
    images_t[:,:,0] = images[0,:,:]
    images_t[:,:,1] = images[1,:,:]
    # import pdb; pdb.set_trace()
    images_t = images_t.astype(np.int8)
    image_raw = images_t.tostring()
    example = tf.train.Example(features=tf.train.Features(feature={
            'height':_int64_feature(rows),
            'width':_int64_feature(cols),
            'depth':_int64_feature(depth),
            'label':_int64_feature(np.where(labels == 1)[0][0]),
            'image_raw':_bytes_feature(image_raw)}))
    index += 1
    # import pdb; pdb.set_trace()
    writer.write(example.SerializeToString())
    if index % 1000 == 0:
        print("%d samples has been processed!" % index)
        if index % 100000 == 0:
            break
writer.close()


# tf.parse_single_example(serialized_example,features={'image_raw': tf.FixedLenFeature([], tf.string),'label': tf.FixedLenFeature([], tf.int64),})
